import matplotlib.pyplot as plt

"""
Lorenz system
dx/dt = -ax + ay
dy/dt = bx -y -xz
dz/dt = -cz + xy
"""
#设定参数
a = 10
b = 28
c = 8/3

# 定义x的微分方程
def fx(x, y, z):
    return -a*x + a*y

# 定义y的微分方程
def fy(x, y, z):
    return b*x - y -x*z

# 定义z的微分方程
def fz(x, y, z):
    return -c*z+x*y

# 利用四阶龙格库塔方法求Lorenz方程
def Lorenz(x0, y0, z0, N):
    h = 0.01
    x = []
    y = []
    z = []
    for i in range(N):
        K1 = fx(x0, y0, z0)
        L1 = fy(x0, y0, z0)
        M1 = fz(x0, y0, z0)

        K2 = fx(x0 + h * K1 / 2, y0 + h * L1 / 2, z0 + h * M1 / 2)
        L2 = fy(x0 + h * K1 / 2, y0 + h * L1 / 2, z0 + h * M1 / 2)
        M2 = fz(x0 + h * K1 / 2, y0 + h * L1 / 2, z0 + h * M1 / 2)

        K3 = fx(x0 + h * K2 / 2, y0 + h * L2 / 2, z0 + h * M2 / 2)
        L3 = fy(x0 + h * K2 / 2, y0 + h * L2 / 2, z0 + h * M2 / 2)
        M3 = fz(x0 + h * K2 / 2, y0 + h * L2 / 2, z0 + h * M2 / 2)

        K4 = fx(x0 + h * K3, y0 + h * L3, z0 + h * M3)
        L4 = fy(x0 + h * K3, y0 + h * L3, z0 + h * M3)
        M4 = fz(x0 + h * K3, y0 + h * L3, z0 + h * M3)
        # x0、y0、z0统一更新
        x0 = x0 + h / 6 * (K1 + 2 * K2 + 2 * K3 + K4)
        y0 = y0 + h / 6 * (L1 + 2 * L2 + 2 * L3 + L4)
        z0 = z0 + h / 6 * (M1 + 2 * M2 + 2 * M3 + M4)

        x.append(x0)
        y.append(y0)
        z.append(z0)
    return x, y, z

def main():
    #迭代次数
    N = 10000
    #设初值
    x0 = 1
    y0 = 0
    z0 = 0
    # 画图
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    x, y, z = Lorenz(x0, y0, z0, N)
    ax.set_xlabel('x(t)')
    ax.set_ylabel('y(t)')
    ax.set_zlabel('z(t)')
    ax.set_title('x0=1 y0=0 z0=0')
    ax.plot(x, y, z, c='g', linewidth=0.5)
    plt.show()


if __name__ == '__main__' :
    main()


